import os
import json
from typing import Dict, List, Any
import requests
import time

# 尝试导入配置文件
try:
    from config import DEEPSEEK_API_KEY as CONFIG_API_KEY
except ImportError:
    CONFIG_API_KEY = None

# 尝试加载.env文件
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

def check_api_key():
    """检查API密钥配置"""
    api_key = os.getenv("DEEPSEEK_API_KEY") or CONFIG_API_KEY
    
    if not api_key:
        print("❌ 请设置DeepSeek API密钥")
        print("\n📝 设置方法：")
        print("方法1: 设置环境变量 DEEPSEEK_API_KEY")
        print("   Windows: set DEEPSEEK_API_KEY=your_api_key_here")
        print("   Linux/Mac: export DEEPSEEK_API_KEY=your_api_key_here")
        print("\n方法2: 在config.py中直接设置DEEPSEEK_API_KEY")
        print("   编辑config.py文件，将DEEPSEEK_API_KEY替换为您的实际密钥")
        print("\n方法3: 创建.env文件并设置DEEPSEEK_API_KEY")
        print("   创建.env文件，添加: DEEPSEEK_API_KEY=your_api_key_here")
        print("\n🔗 获取API密钥: https://platform.deepseek.com/")
        return False
    
    print("✅ API密钥已配置")
    return True

def load_qa_data(file_path: str) -> List[Dict[str, Any]]:
    """加载问答数据"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"❌ 文件 {file_path} 不存在")
        return []
    except json.JSONDecodeError:
        print(f"❌ 文件 {file_path} 格式错误")
        return []

def call_deepseek_api(prompt: str, subject: str, api_key: str) -> Dict[str, Any]:
    """调用DeepSeek API"""
    url = "https://api.deepseek.com/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    # 构建思维链提示词
    system_prompt = f"""你是一个专业的{subject}老师，擅长通过思维链的方式帮助学生理解问题。

请按照以下思维链结构来分析和回答问题：

1. 问题理解 (understanding): 分析问题的核心概念、关键信息和目标
2. 信息检索 (retrieval): 回顾相关的{subject}知识、公式、定理或原理
3. 答案组织 (organization): 整理解题思路、逻辑顺序和推理过程
4. 生成回答 (response): 给出完整、清晰的解答

请确保回答详细、准确，适合学生理解。"""

    user_prompt = f"请为以下{subject}问题生成思维链分析：\n\n{prompt}"
    
    data = {
        "model": "deepseek-chat",
        "messages": [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": user_prompt
            }
        ],
        "temperature": 0.7,
        "max_tokens": 2000
    }
    
    try:
        print(f"   正在调用DeepSeek API...")
        response = requests.post(url, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"❌ API调用失败: {e}")
        return None

def parse_response(response_text: str) -> Dict[str, str]:
    """解析API响应，提取思维链结构"""
    result = {
        "understanding": "",
        "retrieval": "",
        "organization": "",
        "response": ""
    }
    
    # 尝试从响应中提取思维链部分
    lines = response_text.split('\n')
    current_section = None
    
    for line in lines:
        line = line.strip()
        # 检查各种可能的标题格式
        if any(keyword in line.lower() for keyword in ['问题理解', '理解', 'understanding']):
            current_section = 'understanding'
            # 提取冒号后的内容
            if ':' in line or '：' in line:
                result[current_section] = line.split(':', 1)[-1].split('：', 1)[-1].strip()
        elif any(keyword in line.lower() for keyword in ['信息检索', '检索', 'retrieval']):
            current_section = 'retrieval'
            if ':' in line or '：' in line:
                result[current_section] = line.split(':', 1)[-1].split('：', 1)[-1].strip()
        elif any(keyword in line.lower() for keyword in ['答案组织', '组织', 'organization']):
            current_section = 'organization'
            if ':' in line or '：' in line:
                result[current_section] = line.split(':', 1)[-1].split('：', 1)[-1].strip()
        elif any(keyword in line.lower() for keyword in ['生成回答', '回答', 'response', '解答']):
            current_section = 'response'
            if ':' in line or '：' in line:
                result[current_section] = line.split(':', 1)[-1].split('：', 1)[-1].strip()
        elif current_section and line:
            # 如果当前有活动部分且行不为空，添加到该部分
            if result[current_section]:
                result[current_section] += '\n' + line
            else:
                result[current_section] = line
    
    # 如果没有找到结构化内容，将整个响应作为回答
    if not any(result.values()):
        result['response'] = response_text.strip()
    
    return result

def main():
    """主函数"""
    print("🎓 学生问答思维链生成器 (DeepSeek版本)")
    print("=" * 50)
    
    # 检查API密钥
    if not check_api_key():
        return
    
    print("🚀 启动生成器...")
    
    # 加载数据
    qa_data = load_qa_data("student_qa_dialogue.json")
    if not qa_data:
        print("❌ 无法加载问答数据")
        return
    
    print(f"📚 加载了 {len(qa_data)} 个问题")
    
    # 获取API密钥
    api_key = os.getenv("DEEPSEEK_API_KEY") or CONFIG_API_KEY
    
    # 处理每个问题
    for i, qa_item in enumerate(qa_data, 1):
        print(f"\n🔄 处理第 {i}/{len(qa_data)} 个问题: {qa_item['prompt'][:50]}...")
        
        # 调用API
        subject = qa_item.get('subject', '通用')
        api_response = call_deepseek_api(qa_item['prompt'], subject, api_key)
        
        if not api_response:
            print(f"❌ 跳过问题 {i}")
            continue
        
        # 解析响应
        try:
            content = api_response['choices'][0]['message']['content']
            print(f"   API响应长度: {len(content)} 字符")
            
            parsed = parse_response(content)
            
            # 更新数据
            qa_item['thought_chain'] = {
                'understanding': parsed['understanding'],
                'retrieval': parsed['retrieval'],
                'organization': parsed['organization']
            }
            qa_item['response'] = parsed['response']
            
            print(f"✅ 问题 {i} 处理完成")
            print(f"   学科: {subject}")
            print(f"   理解: {parsed['understanding'][:50]}..." if parsed['understanding'] else "   理解: 未提取到")
            print(f"   检索: {parsed['retrieval'][:50]}..." if parsed['retrieval'] else "   检索: 未提取到")
            print(f"   组织: {parsed['organization'][:50]}..." if parsed['organization'] else "   组织: 未提取到")
            print(f"   回答: {parsed['response'][:50]}..." if parsed['response'] else "   回答: 未提取到")
            
        except (KeyError, IndexError) as e:
            print(f"❌ 解析响应失败: {e}")
            continue
        
        # 添加延迟避免API限制
        if i < len(qa_data):
            print("   等待2秒后处理下一个问题...")
            time.sleep(2)
    
    # 保存结果
    try:
        with open("student_qa_processed.json", 'w', encoding='utf-8') as f:
            json.dump(qa_data, f, ensure_ascii=False, indent=2)
        print(f"\n✅ 处理完成！结果已保存到 student_qa_processed.json")
        
        # 显示结果摘要
        print("\n📊 处理结果摘要:")
        success_count = 0
        for i, qa_item in enumerate(qa_data, 1):
            if qa_item['thought_chain']['understanding'] or qa_item['response']:
                success_count += 1
                print(f"问题 {i}: {qa_item['prompt'][:40]}... ✅")
            else:
                print(f"问题 {i}: {qa_item['prompt'][:40]}... ❌")
        
        print(f"\n📈 成功率: {success_count}/{len(qa_data)} ({success_count/len(qa_data)*100:.1f}%)")
            
    except Exception as e:
        print(f"❌ 保存文件失败: {e}")

if __name__ == "__main__":
    main() 